//
// Copyright 2019 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS-IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
//
package sched

import (
	"errors"

	"google.golang.org/grpc/codes"
	"google.golang.org/grpc/status"
	"github.com/google/schedviz/tracedata/trace"
)

// eventLoader provides a flexible interface for generating threadTransitions
// from raw trace.Events.  A new eventLoader may be populated with a number of
// loader functions, each associated with a distinct tracepoint event name.
// Each loader function converts a single event of the specified type into
// zero or more threadTransitions, describing the state or CPU transitions that
// event described on individual threads.  The threadTransitions thus produced
// may have Unknown values which may then be inferred from subsequent passes,
// and reduced confidence in the event types may be reflected in a generated
// threadTransition by signaling that, in a conflict with other
// threadTransitions, this one should be dropped.
type eventLoader struct {
	stringBank *stringBank
	loaders    map[string]func(*trace.Event, *ThreadTransitionSetBuilder) error
}

// newEventLoader returns a new, empty, eventLoader.
func newEventLoader(loaders map[string]func(*trace.Event, *ThreadTransitionSetBuilder) error, stringBank *stringBank) (*eventLoader, error) {
	if len(loaders) == 0 {
		return nil, errors.New("an empty eventLoader cannot generate threadTransitions")
	}
	return &eventLoader{
		stringBank: stringBank,
		loaders:    loaders,
	}, nil
}

// threadTransitions returns the list of threadTransitions generated by the
// receiver's loaders for the provided trace.Event.
func (el *eventLoader) threadTransitions(ev *trace.Event) ([]*threadTransition, error) {
	if ev.Clipped {
		return nil, nil
	}
	ttsb := newThreadTransitionSetBuilder(el.stringBank)
	loader, ok := el.loaders[ev.Name]
	if !ok {
		return nil, nil
	}
	if err := loader(ev, ttsb); err != nil {
		return nil, err
	}
	return ttsb.transitions()
}

// MissingFieldError is used to report a missing field.
func MissingFieldError(fieldName string, ev *trace.Event) error {
	return status.Errorf(codes.NotFound, "field '%s' not found for event %d", fieldName, ev.Index)
}

// MigrateData comprises the data extracted from a raw sched_migrate_task
// event.
type MigrateData struct {
	PID              PID
	Comm             string
	Priority         Priority
	OrigCPU, DestCPU CPUID
}

// LoadMigrateData loads the data from a sched_migrate event, converting all
// fields to suitable types, and returns a MigrateData struct.
func LoadMigrateData(ev *trace.Event) (*MigrateData, error) {
	ret := &MigrateData{}
	pid, ok := ev.NumberProperties["pid"]
	if !ok {
		return nil, MissingFieldError("pid", ev)
	}
	ret.PID = PID(pid)
	ret.Comm = ev.TextProperties["comm"]
	prio, ok := ev.NumberProperties["prio"]
	ret.Priority = Priority(prio)
	if !ok {
		ret.Priority = UnknownPriority
	}
	origCPU, ok := ev.NumberProperties["orig_cpu"]
	if !ok {
		return nil, MissingFieldError("orig_cpu", ev)
	}
	ret.OrigCPU = CPUID(origCPU)
	destCPU, ok := ev.NumberProperties["dest_cpu"]
	if !ok {
		return nil, MissingFieldError("dest_cpu", ev)
	}
	ret.DestCPU = CPUID(destCPU)
	return ret, nil
}

// SwitchData comprises the data extracted from a raw sched_switch event.
// In its members, Next and Prev refer to the switched-in and switched-out
// threads respectively.
type SwitchData struct {
	NextPID, PrevPID           PID
	NextComm, PrevComm         string
	NextPriority, PrevPriority Priority
	PrevState                  ThreadState
}

// LoadSwitchData loads the data from a sched_switch event, converting all
// fields to suitable types, and returns a SwitchData struct.
func LoadSwitchData(ev *trace.Event) (*SwitchData, error) {
	ret := &SwitchData{}
	nextPID, ok := ev.NumberProperties["next_pid"]
	if !ok {
		return nil, MissingFieldError("next_pid", ev)
	}
	ret.NextPID = PID(nextPID)
	ret.NextComm = ev.TextProperties["next_comm"]
	nextPrio, ok := ev.NumberProperties["next_prio"]
	ret.NextPriority = Priority(nextPrio)
	if !ok {
		ret.NextPriority = UnknownPriority
	}
	prevPID, ok := ev.NumberProperties["prev_pid"]
	if !ok {
		return nil, MissingFieldError("prev_pid", ev)
	}
	ret.PrevPID = PID(prevPID)
	ret.PrevComm = ev.TextProperties["prev_comm"]
	prevPrio, ok := ev.NumberProperties["prev_prio"]
	ret.PrevPriority = Priority(prevPrio)
	if !ok {
		ret.PrevPriority = UnknownPriority
	}
	// The new PID's state is assumed to be RUNNING_STATE, and the old PID's task
	// state will reveal whether it's WAITING_STATE (prev_state == 0,
	// TASK_RUNNING, or 256, TASK_REPORT_MAX, which can signify a preemption) or
	// SLEEPING_STATE (otherwise); The possible values of prevTaskState are
	// defined in sched.h in the kernel.
	prevTaskState, ok := ev.NumberProperties["prev_state"]
	if !ok {
		return nil, MissingFieldError("prev_state", ev)
	}
	ret.PrevState = WaitingState
	if prevTaskState != 0 && prevTaskState != 256 {
		ret.PrevState = SleepingState
	}
	return ret, nil
}
